Zigzag iterator

Time: O(N); Space: O(K); medium

Given two 1d vectors, implement an iterator to return their elements alternately.

Example 1:

Input: v1 = [1, 2], v2 = [3, 4, 5, 6]

Output: [1, 3, 2, 4, 5, 6]

Explanation:

  • By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1, 3, 2, 4, 5, 6].

Example 2:

Input: v1 = [1, 1, 1, 1], v2 = [3, 4, 5, 6]

Output: [1, 3, 1, 4, 1, 5, 1, 6]

[10]:

import collections

class ZigzagIterator(object):
    def __init__(self, v1, v2):
        """
        Initialize your q structure here.
        :type v1: List[int]
        :type v2: List[int]
        """
        self.q = collections.deque([(len(v), iter(v)) for v in (v1, v2) if v])

    def next(self):
        """
        :rtype: int
        """
        len, iter = self.q.popleft()
        if len > 1:
            self.q.append((len-1, iter))
        return next(iter)

    def hasNext(self):
        """
        :rtype: bool
        """
        return bool(self.q)

[11]:

v1 = [1, 2]
v2 = [3, 4, 5, 6]
s = ZigzagIterator(v1, v2)
res = [1, 3, 2, 4, 5, 6]
i = 0
while s.hasNext():
    assert s.next() == res[i]
    i +=1


v1 = [1, 1, 1, 1]
v2 = [3, 4, 5, 6]
s = ZigzagIterator(v1, v2)
res =  [1, 3, 1, 4, 1, 5, 1, 6]
i = 0
while s.hasNext():
    assert s.next() == res[i]
    i +=1